import concurrent.futures
import requests
import os
import random
import string
import logging
from openpyxl import Workbook
from time import time

# 设置日志
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')

# 指定生成文件的数量和大小
num_files_to_generate = 200  # 生成文件的数量
size_in_kilobytes = 1024  # 文件大小(以KB为单位，这里设置为1MB)
output_folder = 'filesOutput'  # 输出脚本所在目录文件夹


def generate_files(num_files, size_in_kb, output_folder):
    """生成指定数量和大小的文件并放入指定文件夹"""
    if not os.path.exists(output_folder):
        os.makedirs(output_folder)
    filenames = []
    for _ in range(num_files):
        size_in_bytes = size_in_kb * 1024  # 转换为字节
        filename = ''.join(random.choices(string.ascii_lowercase + string.digits, k=10)) + '.txt'
        filepath = os.path.join(output_folder, filename)
        with open(filepath, 'wb') as f:
            f.write(os.urandom(size_in_bytes))
        filenames.append(filepath)
    return filenames


def send_request_and_delete_file(url, headers, filepath):
    """发送文件以及自定义请求头到url，记录响应码和响应时间"""
    try:
        with open(filepath, 'rb') as f:
            filename = os.path.basename(filepath)
            files = {'file': (filename, f)}
            start_time = time()
            response = requests.post(url, headers=headers, files=files)
            response_time = time() - start_time
            logging.info(f"{filepath} - 响应码: {response.status_code}, 响应时间: {response_time}s")
            return filepath, response.status_code, response_time
    except Exception as e:
        logging.error(f"上传文件发生异常 {filepath}: {e}")


def cleanup_files(filepaths):
    """删除所有生成的文件"""
    for filepath in filepaths:
        try:
            os.remove(filepath)
        except OSError as e:
            logging.warning(f"删除文件发生错误 {filepath}: {e}")


def record_to_excel(records, filename='results.xlsx'):
    """记录结果到excel中."""
    wb = Workbook()
    ws = wb.active
    ws.append(['文件名', '响应码', '响应时间'])
    success_count = 0

    for record in records:
        ws.append(record)
        if record[1] == 200:
            success_count += 1

    success_rate = success_count / len(records) * 100
    ws.append([])
    ws.append(['成功率:', f'{success_rate}%'])
    wb.save(filename)


# 设置并发数和持续时间
concurrent_requests = 50  # 并发数
test_duration_seconds = 1  # 持续时间（秒）

# 自定义请求头
headers = {
    # 'Content-Type': 'multipart/form-data', # 这个注释保留，因为requests会自动添加
    'X-App-Id': 'ndtzOjl7',
    'X-Timestamp': 'ndtzOjl7',
    'X-Sequence-No': 'ndtzOjl7',
    'X-Signature': 'ndtzOjl7'
}

# 生成文件
file_paths = generate_files(num_files_to_generate, size_in_kilobytes, output_folder)

# 开始测试
start_time = time()
records = []

url = 'http://10.10.2.236:32287/pay/upload'  # 上传文件地址

with concurrent.futures.ThreadPoolExecutor(max_workers=concurrent_requests) as executor:
    futures = [executor.submit(send_request_and_delete_file, url, headers, filepath) for filepath in file_paths]

    # 等待所有的futures完成
    for future in concurrent.futures.as_completed(futures):
        result = future.result()
        if result:
            records.append(result)

# 测试结束，记录结果到Excel
record_to_excel(records)

# 清理生成的文件
cleanup_files(file_paths)
logging.info('测速和结束，结果生成在 results.xlsx.')
